import csv
import matplotlib.pyplot as pyplot
import math
import Sugarscape

class Plotting:
	
	def __init__(self, f):
		self.csv = csv.reader(open(f))
		self.enumeration = {"tax_rate" : 0,
							"standard deviation" : 1,
							"average" : 2,
							"median" : 3,
							"maximum" : 4,
							"minimum" : 5,
							"range" : 6,
							"total_wealth" : 7,
							"gini" : 8,
                            "bottom_quartile" : 9
							}
	
	def plot_arbitrary(self, field1, field2, xlog=False, ylog=False):
		int1 = self.enumeration[field1]
		int2 = self.enumeration[field2]
		keys, vals = [],[]
		
		for line in self.csv:
			if line[0] == "tax_rate": continue
			keys.append(float(line[int1]) + .001)
			vals.append(float(line[int2]) + .001)
		pyplot.plot(keys, vals,'o')
		pyplot.xlabel(field1)
		pyplot.ylabel(field2)
		if xlog:
			pyplot.xscale('log')
		if ylog:
			pyplot.yscale('log')
		pyplot.title("%s vs %s" %(field1, field2))
		pyplot.show()
	
	def plot_ratio(self, field1, numerator_field, denominator_field, xlog = False, ylog = False, xlog_num = False, ylog_num = False):
		int1 = self.enumeration[field1]
		int2 = self.enumeration[numerator_field]
		int3 = self.enumeration[denominator_field]
		
		keys, vals = [],[]
		
		for line in self.csv:
			if line[0] == "tax_rate": continue
			keys.append(float(line[int1]))
			num = (float(line[int2]))
			den = (float(line[int3]))
			if xlog_num:
				num = math.log(num)
			if ylog_num:
				den = math.log(den)
			vals.append(num / float(den))
		print keys
		pyplot.plot(keys, vals,'o')
		pyplot.xlabel(field1)
		pyplot.ylabel(numerator_field + " over " + denominator_field)
		if xlog:
			pyplot.xscale('log')
		if ylog:
			pyplot.yscale('log')
		pyplot.title(numerator_field + " over " + denominator_field + " vs " + field1)
		pyplot.show()	


	def lorenz(self, tax_rate):
		m = Sugarscape.Sugarscape(tax_rate=tax_rate)
		for i in range(100):
			m.nextstep()
		datafile = csv.writer(open('lorenz.csv', 'w'))
		agent_wealths = []
		for agent in m.agents:
			agent_wealths.append(agent.sugar_reserve)
			datafile.writerow((agent.sugar_reserve,0))
		
		agent_wealths.sort()
		total_wealth = sum(agent_wealths)
		agent_props = []
		
		agents = []
		for i in range(100):
			 agent_props.append(sum([agent_wealths[j] for j in range((len(agent_wealths)*i)/100)]) / float(total_wealth))
		return agent_props

	def lorenz_plot(self):
		#hard-coded
		tax_style = [(0,'o-'), (.05,'s-'), (.25,'d-')]
		
		for tax,style in tax_style:
			pyplot.plot([x for x in range(100)], [x*100 for x in self.lorenz(tax)],style,label='%g%s tax rate' %(tax*100,"%"),linewidth=1)
			pyplot.hold(True)
		pyplot.plot([x for x in range(100)], [x for x in range(100)], label='perfect equality', linewidth=5)
		pyplot.legend(loc=2)
		pyplot.grid(which='major')
		pyplot.xlabel('percent of population', fontsize=15)
		pyplot.ylabel('percent of wealth', fontsize=15)
		pyplot.title('Lorenz Curves for Different Sugarscape Tax Rates',fontsize=20)
		pyplot.show()
	
	def get_fields(self, name1, name2):
		int1 = self.enumeration[name1]
		int2 = self.enumeration[name2]
		
		retlist1 = []
		retlist2 = []
		for line in self.csv:
			if line[0]=="tax_rate": continue
			retlist1.append(float(line[int1]))
			retlist2.append(float(line[int2]))
		return retlist1, retlist2
	
	def tax_rate_bottom(self):
		taxes, bottom = self.get_fields("tax_rate", "bottom_quartile")
		pyplot.plot(taxes, bottom,'bo')
		pyplot.grid(which='major')
		pyplot.xlabel('tax rate, percent', fontsize=15)
		pyplot.ylabel('average wealth of bottom quartile', fontsize=15)
		pyplot.axis([0, max(taxes), 0, max(bottom)+1])
		pyplot.title('Mean wealth, bottom quartile',fontsize=20)
		pyplot.show()
	
	def tax_rate_gini(self):
		taxes, gini = self.get_fields("tax_rate", "gini")
		pyplot.plot(taxes, gini,'bo')
		pyplot.grid(which='major')
		pyplot.xlabel('tax rate, percent', fontsize=15)
		pyplot.ylabel('gini coefficient', fontsize=15)
		pyplot.title('Gini coefficient vs tax rate',fontsize=20)
		pyplot.show()
	
	def median_gini(self):
		median, gini = self.get_fields("median", "gini")
		pyplot.plot(gini, median,'bo')
		pyplot.grid(which='major')
		pyplot.ylabel('median', fontsize=15)
		pyplot.xlabel('gini coefficient', fontsize=15)
		pyplot.title('median versus gini coefficient',fontsize=20)
		pyplot.show()
	
	def tax_rate_mean(self):
		mean, tax_rate = self.get_fields("average", "tax_rate")
		pyplot.plot(tax_rate, mean,'bo')
		pyplot.grid(which='major')
		pyplot.ylabel('mean', fontsize=15)
		pyplot.xlabel('tax rate, percent', fontsize=15)
		pyplot.axis([0, max(tax_rate), 0, max(mean)+5])
		pyplot.title('Mean wealth vs tax rate',fontsize=20)
		pyplot.show()
	
	def __str__(self):
		retstr = ""
		for line in self.csv:
			for val in line:
				retstr = retstr + str(val) + "\t"
			retstr = retstr + "\n"
		return retstr
	

if __name__=="__main__":
	m= Plotting('tax_data_changedparam.csv')
	#~ m.lorenz_plot()
	m.tax_rate_bottom()
	#m.tax_rate_mean()
	#m.tax_rate_gini()
	
	#~ m.plot_arbitrary("tax_rate", "total_wealth")
	
	#~ print m
	#~ m.plot_arbitrary('tax_rate','minimum') #,ylog=True, xlog=True)
	#~ m.plot_arbitrary('tax_rate','bottom') #,ylog=True, xlog=True)
	#~ m.plot_arbitrary('tax_rate','gini') #,ylog=True, xlog=True)
	#~ m.plot_arbitrary('gini','total_wealth', xlog=True) #, xlog=True)
	#~ m.plot_ratio("tax_rate","total_wealth","standard deviation",ylog=True)
	
